#!/usr/bin/env python
# coding=utf-8
import math
from typing import List, Tuple, Dict, Optional, Any

import torch
import torchvision
from torch import nn, Tensor

from .image_list import ImageList


def _resize_image(image, self_min_size, self_max_size):
    # type: (Tensor, float, float) -> Tensor
    im_shape = torch.tensor(image.shape[-2:])
    min_size = float(torch.min(im_shape))    # 获取高宽中的最小值
    max_size = float(torch.max(im_shape))    # 获取高宽中的最大值
    scale_factor = self_min_size / min_size  # 根据指定最小边长和图片最小边长计算缩放比例

    # 如果使用该缩放比例计算的图片最大边长大于指定的最大边长,则将缩放比例设为指定最大边长和图片最大边长之比
    if max_size * scale_factor > self_max_size:
        scale_factor = self_max_size / max_size   

    # interpolate利用插值的方法缩放图片
    # image[None]操作是在最前面添加batch维度[C, H, W] -> [1, C, H, W]
    # bilinear只支持4D Tensor
    # 最后用[0]索引变回[C, H, W]
    image = torch.nn.functional.interpolate(
        image[None], scale_factor=scale_factor, mode="bilinear", recompute_scale_factor=True,
        align_corners=False)[0]

    return image


def resize_boxes(boxes, original_size, new_size):
    # type: (Tensor, List[int], List[int]) -> Tensor
    """
    将boxes参数根据图像的缩放情况进行相应缩放

    Arguments:
        original_size: 图像缩放前的尺寸
        new_size: 图像缩放后的尺寸
    """
    ratios = [
        torch.tensor(s, dtype=torch.float32, device=boxes.device) /
        torch.tensor(s_orig, dtype=torch.float32, device=boxes.device)
        for s, s_orig in zip(new_size, original_size)
    ]
    ratios_height, ratios_width = ratios
    # Removes a tensor dimension, boxes [minibatch, 4]
    # Returns a tuple of all slices along a given dimension, already without it.
    xmin, ymin, xmax, ymax = boxes.unbind(1)
    xmin = xmin * ratios_width
    xmax = xmax * ratios_width
    ymin = ymin * ratios_height
    ymax = ymax * ratios_height
    return torch.stack((xmin, ymin, xmax, ymax), dim=1)



# class GeneralizedRCNNTransform(nn.Module):
#     """
#     Performs input / target transformation before feeding the data to a GeneralizedRCNN
#     model.

#     The transformations it perform are:
#         - input normalization (mean subtraction and std division)
#         - input / target resizing to match min_size / max_size

#     It returns a ImageList for the inputs, and a List[Dict[Tensor]] for the targets
#     """

#     def __init__(self, min_size, max_size, image_mean, image_std):
#         super(GeneralizedRCNNTransform, self).__init__()
#         if not isinstance(min_size, (list, tuple)):
#             min_size = (min_size,)
#         self.min_size = min_size      # 指定图像的最小边长范围
#         self.max_size = max_size      # 指定图像的最大边长范围
#         self.image_mean = image_mean  # 指定图像在标准化处理中的均值
#         self.image_std = image_std    # 指定图像在标准化处理中的方差

#     def normalize(self, image):
#         """标准化处理"""
#         dtype, device = image.dtype, image.device
#         mean = torch.as_tensor(self.image_mean, dtype=dtype, device=device)
#         std = torch.as_tensor(self.image_std, dtype=dtype, device=device)
#         # [:, None, None]: shape [3] -> [3, 1, 1]
#         return (image - mean[:, None, None]) / std[:, None, None]

#     def torch_choice(self, k):
#         # type: (List[int]) -> int
#         """
#         Implements `random.choice` via torch ops so it can be compiled with
#         TorchScript. Remove if https://github.com/pytorch/pytorch/issues/25803
#         is fixed.
#         torch.empty返回指定size空张量,uniform将使用此方法的tensor对象填充,
#         使用from到to的均匀分布抽样值(此处为0.到float(k长度))
#         """
#         index = int(torch.empty(1).uniform_(0., float(len(k))).item())
#         return k[index]

#     def resize(self, image, target):
#         # type: (Tensor, Optional[Dict[str, Tensor]]) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]
#         """
#         将图片缩放到指定的大小范围内,并对应缩放bboxes信息
#         Args:
#             image: 输入的图片
#             target: 输入图片的相关信息(包括bboxes信息)

#         Returns:
#             image: 缩放后的图片
#             target: 缩放bboxes后的图片相关信息
#         """
#         # image shape is [channel, height, width]
#         h, w = image.shape[-2:]

#         if self.training:
#             size = float(self.torch_choice(self.min_size))  # 指定输入图片的最小边长,注意是self.min_size不是min_size
#         else:
#             # FIXME assume for now that testing uses the largest scale
#             size = float(self.min_size[-1])    # 指定输入图片的最小边长,注意是self.min_size不是min_size

#         image = _resize_image(image, size, float(self.max_size))

#         if target is None:
#             return image, target

#         bbox = target["boxes"]
#         # 根据图像的缩放比例来缩放bbox
#         bbox = resize_boxes(bbox, [h, w], image.shape[-2:])
#         target["boxes"] = bbox

#         return image, target


#     def max_by_axis(self, the_list):
#         # type: (List[List[int]]) -> List[int]
#         maxes = the_list[0]
#         for sublist in the_list[1:]:
#             for index, item in enumerate(sublist):
#                 maxes[index] = max(maxes[index], item)
#         return maxes

#     def batch_images(self, images, size_divisible=32):
#         # type: (List[Tensor], int) -> Tensor
#         """
#         将一批图像打包成一个batch返回（注意batch中每个tensor的shape是相同的）
#         Args:
#             images: 输入的一批图片
#             size_divisible: 将图像高和宽调整到该数的整数倍

#         Returns:
#             batched_imgs: 打包成一个batch后的tensor数据
#         """

#         # 分别计算一个batch中所有图片中的最大channel, height, width
#         max_size = self.max_by_axis([list(img.shape) for img in images])

#         stride = float(size_divisible)
#         # max_size = list(max_size)
#         # 将height向上调整到stride的整数倍
#         max_size[1] = int(math.ceil(float(max_size[1]) / stride) * stride)
#         # 将width向上调整到stride的整数倍
#         max_size[2] = int(math.ceil(float(max_size[2]) / stride) * stride)

#         # [batch, channel, height, width]
#         batch_shape = [len(images)] + max_size

#         # 创建shape为batch_shape且值全部为0的tensor
#         # new_full返回size指定的shape,以fill_value填充,而dtype与device与使用方法的对象(images)相同
#         batched_imgs = images[0].new_full(size=batch_shape, fill_value=0)
#         for img, pad_img in zip(images, batched_imgs):
#             # 将输入images中的每张图片复制到新的batched_imgs的每张图片中，对齐左上角，保证bboxes的坐标不变
#             # 这样保证输入到网络中一个batch的每张图片的shape相同
#             # copy_: Copies the elements from src into self tensor and returns self
#             pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)

#         return batched_imgs

#     def postprocess(self,
#                     result,                # type: List[Dict[str, Tensor]]
#                     image_shapes,          # type: List[Tuple[int, int]]
#                     original_image_sizes   # type: List[Tuple[int, int]]
#                     ):
#         # type: (...) -> List[Dict[str, Tensor]]
#         """
#         对网络的预测结果进行后处理（主要将bboxes还原到原图像尺度上）
#         Args:
#             result: list(dict), 网络的预测结果, len(result) == batch_size
#             image_shapes: list(torch.Size), 图像预处理缩放后的尺寸, len(image_shapes) == batch_size
#             original_image_sizes: list(torch.Size), 图像的原始尺寸, len(original_image_sizes) == batch_size

#         Returns:

#         """
#         if self.training:
#             return result

#         # 遍历每张图片的预测信息，将boxes信息还原回原尺度
#         for i, (pred, im_s, o_im_s) in enumerate(zip(result, image_shapes, original_image_sizes)):
#             boxes = pred["boxes"]
#             boxes = resize_boxes(boxes, im_s, o_im_s)  # 将bboxes缩放回原图像尺度上
#             result[i]["boxes"] = boxes
#         return result

#     def __repr__(self):
#         """自定义输出实例化对象的信息,可通过print打印实例信息"""
#         format_string = self.__class__.__name__ + '('
#         _indent = '\n    '
#         format_string += "{0}Normalize(mean={1}, std={2})".format(_indent, self.image_mean, self.image_std)
#         format_string += "{0}Resize(min_size={1}, max_size={2}, mode='bilinear')".format(_indent, self.min_size,
#                                                                                          self.max_size)
#         format_string += '\n)'
#         return format_string

#     def forward(self,
#                 images,       # type: List[Tensor]
#                 targets=None  # type: Optional[List[Dict[str, Tensor]]]
#                 ):
#         # type: (...) -> Tuple[ImageList, Optional[List[Dict[str, Tensor]]]]
#         images = [img for img in images]
#         for i in range(len(images)):
#             image = images[i]
#             target_index = targets[i] if targets is not None else None

#             if image.dim() != 3:
#                 raise ValueError("images is expected to be a list of 3d tensors "
#                                  "of shape [C, H, W], got {}".format(image.shape))
#             image = self.normalize(image)                # 对图像进行标准化处理
#             image, target_index = self.resize(image, target_index)   # 对图像和对应的bboxes缩放到指定范围
#             images[i] = image
#             if targets is not None and target_index is not None:
#                 targets[i] = target_index

#         # 记录resize后的图像尺寸
#         image_sizes = [img.shape[-2:] for img in images]
#         images = self.batch_images(images)  # 将images打包成一个batch
#         image_sizes_list = torch.jit.annotate(List[Tuple[int, int]], [])

#         for image_size in image_sizes:
#             assert len(image_size) == 2
#             image_sizes_list.append((image_size[0], image_size[1]))
#         # 此处实例化ImageList类
#         # images[2,3,H,W]  
#         # image_sizes 长度为2的List[(h,w),...] 
#         image_list = ImageList(images, image_sizes_list)
#         return image_list, targets


@torch.jit.unused
def _get_shape_onnx(image: Tensor) -> Tensor:
    from torch.onnx import operators

    return operators.shape_as_tensor(image)[-2:]


@torch.jit.unused
def _fake_cast_onnx(v: Tensor) -> float:
    # ONNX requires a tensor but here we fake its type for JIT.
    return v


def _resize_image_and_masks(
    image: Tensor,
    self_min_size: int,
    self_max_size: int,
    target: Optional[Dict[str, Tensor]] = None,
    fixed_size: Optional[Tuple[int, int]] = None,
) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
    if torchvision._is_tracing():
        im_shape = _get_shape_onnx(image)
    elif torch.jit.is_scripting():
        im_shape = torch.tensor(image.shape[-2:])
    else:
        im_shape = image.shape[-2:]

    size: Optional[List[int]] = None
    scale_factor: Optional[float] = None
    recompute_scale_factor: Optional[bool] = None
    if fixed_size is not None:
        size = [fixed_size[1], fixed_size[0]]
    else:
        if torch.jit.is_scripting() or torchvision._is_tracing():
            min_size = torch.min(im_shape).to(dtype=torch.float32)
            max_size = torch.max(im_shape).to(dtype=torch.float32)
            self_min_size_f = float(self_min_size)
            self_max_size_f = float(self_max_size)
            scale = torch.min(self_min_size_f / min_size, self_max_size_f / max_size)

            if torchvision._is_tracing():
                scale_factor = _fake_cast_onnx(scale)
            else:
                scale_factor = scale.item()

        else:
            # Do it the normal way
            min_size = min(im_shape)
            max_size = max(im_shape)
            scale_factor = min(self_min_size / min_size, self_max_size / max_size)

        recompute_scale_factor = True

    image = torch.nn.functional.interpolate(
        image[None],
        size=size,
        scale_factor=scale_factor,
        mode="bilinear",
        recompute_scale_factor=recompute_scale_factor,
        align_corners=False,
    )[0]

    if target is None:
        return image, target

    if "masks" in target:
        mask = target["masks"]
        mask = torch.nn.functional.interpolate(
            mask[:, None].float(), size=size, scale_factor=scale_factor, recompute_scale_factor=recompute_scale_factor
        )[:, 0].byte()
        target["masks"] = mask
    return image, target


class GeneralizedRCNNTransform(nn.Module):
    """
    Performs input / target transformation before feeding the data to a GeneralizedRCNN
    model.

    The transformations it performs are:
        - input normalization (mean subtraction and std division)
        - input / target resizing to match min_size / max_size

    It returns a ImageList for the inputs, and a List[Dict[Tensor]] for the targets
    """

    def __init__(
        self,
        min_size: int,
        max_size: int,
        image_mean: List[float],
        image_std: List[float],
        size_divisible: int = 32,
        fixed_size: Optional[Tuple[int, int]] = None,
        **kwargs: Any,
    ):
        super().__init__()
        if not isinstance(min_size, (list, tuple)):
            min_size = (min_size,)
        self.min_size = min_size
        self.max_size = max_size
        self.image_mean = image_mean
        self.image_std = image_std
        self.size_divisible = size_divisible
        self.fixed_size = fixed_size
        self._skip_resize = kwargs.pop("_skip_resize", False)

    def forward(
        self, images: List[Tensor], targets: Optional[List[Dict[str, Tensor]]] = None
    ) -> Tuple[ImageList, Optional[List[Dict[str, Tensor]]]]:
        images = [img for img in images]
        if targets is not None:
            # make a copy of targets to avoid modifying it in-place
            # once torchscript supports dict comprehension
            # this can be simplified as follows
            # targets = [{k: v for k,v in t.items()} for t in targets]
            targets_copy: List[Dict[str, Tensor]] = []
            for t in targets:
                data: Dict[str, Tensor] = {}
                for k, v in t.items():
                    data[k] = v
                targets_copy.append(data)
            targets = targets_copy
        for i in range(len(images)):
            image = images[i]
            target_index = targets[i] if targets is not None else None

            if image.dim() != 3:
                raise ValueError(f"images is expected to be a list of 3d tensors of shape [C, H, W], got {image.shape}")
            image = self.normalize(image)
            image, target_index = self.resize(image, target_index)
            images[i] = image
            if targets is not None and target_index is not None:
                targets[i] = target_index

        image_sizes = [img.shape[-2:] for img in images]
        images = self.batch_images(images, size_divisible=self.size_divisible)
        image_sizes_list: List[Tuple[int, int]] = []
        for image_size in image_sizes:
            torch._assert(
                len(image_size) == 2,
                f"Input tensors expected to have in the last two elements H and W, instead got {image_size}",
            )
            image_sizes_list.append((image_size[0], image_size[1]))

        image_list = ImageList(images, image_sizes_list)
        return image_list, targets

    def normalize(self, image: Tensor) -> Tensor:
        if not image.is_floating_point():
            raise TypeError(
                f"Expected input images to be of floating type (in range [0, 1]), "
                f"but found type {image.dtype} instead"
            )
        dtype, device = image.dtype, image.device
        mean = torch.as_tensor(self.image_mean, dtype=dtype, device=device)
        std = torch.as_tensor(self.image_std, dtype=dtype, device=device)
        return (image - mean[:, None, None]) / std[:, None, None]

    def torch_choice(self, k: List[int]) -> int:
        """
        Implements `random.choice` via torch ops, so it can be compiled with
        TorchScript and we use PyTorch's RNG (not native RNG)
        """
        index = int(torch.empty(1).uniform_(0.0, float(len(k))).item())
        return k[index]

    def resize(
        self,
        image: Tensor,
        target: Optional[Dict[str, Tensor]] = None,
    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
        h, w = image.shape[-2:]
        if self.training:
            if self._skip_resize:
                return image, target
            size = self.torch_choice(self.min_size)
        else:
            size = self.min_size[-1]
        image, target = _resize_image_and_masks(image, size, self.max_size, target, self.fixed_size)

        if target is None:
            return image, target

        bbox = target["boxes"]
        bbox = resize_boxes(bbox, (h, w), image.shape[-2:])
        target["boxes"] = bbox

        if "keypoints" in target:
            keypoints = target["keypoints"]
            keypoints = resize_keypoints(keypoints, (h, w), image.shape[-2:])
            target["keypoints"] = keypoints
        return image, target

    # _onnx_batch_images() is an implementation of
    # batch_images() that is supported by ONNX tracing.
    @torch.jit.unused
    def _onnx_batch_images(self, images: List[Tensor], size_divisible: int = 32) -> Tensor:
        max_size = []
        for i in range(images[0].dim()):
            max_size_i = torch.max(torch.stack([img.shape[i] for img in images]).to(torch.float32)).to(torch.int64)
            max_size.append(max_size_i)
        stride = size_divisible
        max_size[1] = (torch.ceil((max_size[1].to(torch.float32)) / stride) * stride).to(torch.int64)
        max_size[2] = (torch.ceil((max_size[2].to(torch.float32)) / stride) * stride).to(torch.int64)
        max_size = tuple(max_size)

        # work around for
        # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
        # which is not yet supported in onnx
        padded_imgs = []
        for img in images:
            padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]
            padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))
            padded_imgs.append(padded_img)

        return torch.stack(padded_imgs)

    def max_by_axis(self, the_list: List[List[int]]) -> List[int]:
        maxes = the_list[0]
        for sublist in the_list[1:]:
            for index, item in enumerate(sublist):
                maxes[index] = max(maxes[index], item)
        return maxes

    def batch_images(self, images: List[Tensor], size_divisible: int = 32) -> Tensor:
        if torchvision._is_tracing():
            # batch_images() does not export well to ONNX
            # call _onnx_batch_images() instead
            return self._onnx_batch_images(images, size_divisible)

        max_size = self.max_by_axis([list(img.shape) for img in images])
        stride = float(size_divisible)
        max_size = list(max_size)
        max_size[1] = int(math.ceil(float(max_size[1]) / stride) * stride)
        max_size[2] = int(math.ceil(float(max_size[2]) / stride) * stride)

        batch_shape = [len(images)] + max_size
        batched_imgs = images[0].new_full(batch_shape, 0)
        for i in range(batched_imgs.shape[0]):
            img = images[i]
            batched_imgs[i, : img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)

        return batched_imgs

    def postprocess(
        self,
        result: List[Dict[str, Tensor]],
        image_shapes: List[Tuple[int, int]],
        original_image_sizes: List[Tuple[int, int]],
    ) -> List[Dict[str, Tensor]]:
        if self.training:
            return result
        for i, (pred, im_s, o_im_s) in enumerate(zip(result, image_shapes, original_image_sizes)):
            boxes = pred["boxes"]
            boxes = resize_boxes(boxes, im_s, o_im_s)
            result[i]["boxes"] = boxes
            # if "masks" in pred:
            #     masks = pred["masks"]
            #     masks = paste_masks_in_image(masks, boxes, o_im_s)
            #     result[i]["masks"] = masks
            if "keypoints" in pred:
                keypoints = pred["keypoints"]
                keypoints = resize_keypoints(keypoints, im_s, o_im_s)
                result[i]["keypoints"] = keypoints
        return result

    def __repr__(self) -> str:
        format_string = f"{self.__class__.__name__}("
        _indent = "\n    "
        format_string += f"{_indent}Normalize(mean={self.image_mean}, std={self.image_std})"
        format_string += f"{_indent}Resize(min_size={self.min_size}, max_size={self.max_size}, mode='bilinear')"
        format_string += "\n)"
        return format_string


def resize_keypoints(keypoints: Tensor, original_size: List[int], new_size: List[int]) -> Tensor:
    ratios = [
        torch.tensor(s, dtype=torch.float32, device=keypoints.device)
        / torch.tensor(s_orig, dtype=torch.float32, device=keypoints.device)
        for s, s_orig in zip(new_size, original_size)
    ]
    ratio_h, ratio_w = ratios
    resized_data = keypoints.clone()
    if torch._C._get_tracing_state():
        resized_data_0 = resized_data[:, :, 0] * ratio_w
        resized_data_1 = resized_data[:, :, 1] * ratio_h
        resized_data = torch.stack((resized_data_0, resized_data_1, resized_data[:, :, 2]), dim=2)
    else:
        resized_data[..., 0] *= ratio_w
        resized_data[..., 1] *= ratio_h
    return resized_data


def resize_boxes(boxes: Tensor, original_size: List[int], new_size: List[int]) -> Tensor:
    ratios = [
        torch.tensor(s, dtype=torch.float32, device=boxes.device)
        / torch.tensor(s_orig, dtype=torch.float32, device=boxes.device)
        for s, s_orig in zip(new_size, original_size)
    ]
    ratio_height, ratio_width = ratios
    xmin, ymin, xmax, ymax = boxes.unbind(1)

    xmin = xmin * ratio_width
    xmax = xmax * ratio_width
    ymin = ymin * ratio_height
    ymax = ymax * ratio_height
    return torch.stack((xmin, ymin, xmax, ymax), dim=1)





